import copy
from datetime import datetime
from typing import Optional, Sequence, Union
import warnings

from mmengine.config import Config
from mmengine.registry import DATASETS, VISUALIZERS, init_default_scope
from mmengine.visualization import Visualizer
from mmdet.apis import DetInferencer as MmdetInferencer
from mmdet.evaluation import get_classes
from mmdet.utils import ConfigType
import numpy as np

from ..utils import DataHelper
from .base_inferencer import BaseInferencer


class DetInferencer(BaseInferencer, MmdetInferencer):
    def __init__(self,
                 config: str = None,
                 om_model: str = None,
                 device: int = 0) -> None:
        # A global counter tracking the number of images processed, for
        # naming of the output images
        self.num_visualized_imgs = 0
        self.num_predicted_imgs = 0
        init_default_scope('mmdet')

        self.scope = 'mmdet'
        cfg = Config.fromfile(config)

        BaseInferencer.__init__(self, cfg, om_model, device)
        self.dataset_meta = self._init_dataset_meta(cfg)
        self.visualizer = self._init_visualizer(cfg)

    def _init_dataset_meta(self, cfg):
        dataset_meta = {'classes': get_classes('coco')}
        test_dataset_cfg = copy.deepcopy(cfg.test_dataloader.dataset)
        # lazy init. We only need the metainfo.
        test_dataset_cfg['lazy_init'] = True
        metainfo = DATASETS.build(test_dataset_cfg).metainfo
        cfg_palette = metainfo.get('palette', None)
        if cfg_palette is not None:
            dataset_meta['palette'] = cfg_palette
        else:
            if 'palette' not in dataset_meta:
                warnings.warn(
                    'palette does not exist, random is used by default. '
                    'You can also set the palette to customize.')
                dataset_meta['palette'] = 'random'
        return dataset_meta

    def _init_visualizer(self, cfg: ConfigType) -> Optional[Visualizer]:
        """Initialize visualizers.

        Args:
            cfg (ConfigType): Config containing the visualizer information.

        Returns:
            Visualizer or None: Visualizer initialized with config.
        """
        if 'visualizer' not in cfg:
            return None
        timestamp = str(datetime.timestamp(datetime.now()))
        name = cfg.visualizer.get('name', timestamp)
        if Visualizer.check_instance_created(name):
            name = f'{name}-{timestamp}'
        cfg.visualizer.name = name
        visualizer = VISUALIZERS.build(cfg.visualizer)
        visualizer.dataset_meta = self.dataset_meta
        return visualizer

    def __call__(self,
                 inputs: Union[str, Sequence[str]],
                 return_vis: bool = False,
                 show: bool = False,
                 wait_time: int = 0,
                 no_save_vis: bool = False,
                 draw_pred: bool = True,
                 pred_score_thr: float = 0.3,
                 return_datasample: bool = False,
                 print_result: bool = False,
                 no_save_pred: bool = True,
                 out_dir: str = '',
                 **kwargs) -> dict:
        """Call the inferencer.

        Args:
            inputs (Union[str, Sequence[str]]): Inputs for the inferencer.
            show (bool): Whether to display the visualization results in a
                popup window. Defaults to False.
            wait_time (float): The interval of show (s). Defaults to 0.
            no_save_vis (bool): Whether to force not to save prediction
                vis results. Defaults to False.
            draw_pred (bool): Whether to draw predicted bounding boxes.
                Defaults to True.
            pred_score_thr (float): Minimum score of bboxes to draw.
                Defaults to 0.3.
            return_datasample (bool): Whether to return results as
                :obj:`DetDataSample`. Defaults to False.
            print_result (bool): Whether to print the inference result w/o
                visualization to the console. Defaults to False.
            no_save_pred (bool): Whether to force not to save prediction
                results. Defaults to True.
            out_file: Dir to save the inference results or
                visualization. If left as empty, no file will be saved.
                Defaults to ''.

            **kwargs: Other keyword arguments passed to :meth:`preprocess`,
                :meth:`forward`, :meth:`visualize` and :meth:`postprocess`.
                Each key in kwargs should be in the corresponding set of
                ``preprocess_kwargs``, ``forward_kwargs``, ``visualize_kwargs``
                and ``postprocess_kwargs``.

        Returns:
            dict: Inference and visualization results.
        """
        (
            preprocess_kwargs,
            forward_kwargs,
            visualize_kwargs,
            postprocess_kwargs,
        ) = self._dispatch_kwargs(**kwargs)

        data_helper = DataHelper(inputs, self.batch_size)
        results_dict = {'predictions': [], 'visualization': []}
        for batch_images in data_helper.batch_list:
            preds = self.predict(batch_images)

            visualization = self.visualize(
                batch_images,
                preds,
                return_vis=return_vis,
                show=show,
                wait_time=wait_time,
                draw_pred=draw_pred,
                pred_score_thr=pred_score_thr,
                no_save_vis=no_save_vis,
                img_out_dir=out_dir,
                **visualize_kwargs)
            results = self.postprocess(
                preds,
                visualization,
                return_datasample=return_datasample,
                print_result=print_result,
                no_save_pred=no_save_pred,
                pred_out_dir=out_dir,
                **postprocess_kwargs)
            results_dict['predictions'].extend(results['predictions'])
            if results['visualization'] is not None:
                results_dict['visualization'].extend(results['visualization'])
        return results_dict
